import pybaobabdt
import pandas as pd
from scipy.io import arff
from sklearn.tree import DecisionTreeClassifier
data = arff.loadarff('data/winequality-red.arff')
df = pd.DataFrame(data[0])
y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]
clf = DecisionTreeClassifier().fit(X,y);
ax = pybaobabdt.drawTree(clf, size=10, dpi=72, features=features)
<Figure size 720x720 with 0 Axes>
ax.get_figure().savefig('visualizations/tree.png', format='png', dpi=300, transparent=True)
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
data = arff.loadarff('data/vehicle.arff')
df = pd.DataFrame(data[0])
y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]
clf = RandomForestClassifier(n_estimators=20, n_jobs=-1, random_state=0);
clf.fit(X, y);
size = (15,15)
plt.rcParams['figure.figsize'] = size
fig = plt.figure(figsize=size, dpi=300)
for idx, tree in enumerate(clf.estimators_):
ax1 = fig.add_subplot(5, 4, idx+1)
pybaobabdt.drawTree(tree, model=clf, size=15, dpi=300, features=features, ax=ax1)
fig.savefig('visualizations/random-forest.png', format='png', dpi=1200, transparent=True)